#include "sqlformatter.h"
#include <QRegularExpression>
#include <QRegularExpressionMatch>
#include <QRegularExpressionMatchIterator>
#include <QStringList>
#include <QSet>

SQLFormatter::SQLFormatter()
    : m_indentSize(4)
    , m_uppercaseKeywords(true)
{
    // 初始化SQL关键字列表
    QStringList keywordList = {
        "SELECT", "FROM", "WHERE", "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER", 
        "TABLE", "INDEX", "VIEW", "TRIGGER", "PROCEDURE", "FUNCTION", "DATABASE", "SCHEMA",
        "PRIMARY", "KEY", "FOREIGN", "REFERENCES", "CONSTRAINT", "UNIQUE", "CHECK", "DEFAULT",
        "NOT", "NULL", "AND", "OR", "IN", "EXISTS", "BETWEEN", "LIKE", "AS", "ON", "JOIN", 
        "INNER", "LEFT", "RIGHT", "OUTER", "CROSS", "FULL", "UNION", "ALL", "DISTINCT", 
        "GROUP", "BY", "ORDER", "HAVING", "LIMIT", "OFFSET", "INTO", "VALUES", "SET", 
        "ADD", "COLUMN", "TO", "IF", "ELSE", "BEGIN", "END", "WHEN", "THEN", "CASE", 
        "WITH", "RECURSIVE", "OVER", "PARTITION", "ROW", "ROWS", "RANGE", "PRECEDING", 
        "FOLLOWING", "CURRENT", "FOR", "ONLY", "USING", "RETURNING", "CASCADE", "RESTRICT"
    };
    
    for (const QString &keyword : keywordList) {
        m_keywords[keyword.toUpper()] = true;
    }
}

QString SQLFormatter::format(const QString &sql) const
{
    // 对SQL进行词法分析
    QList<Token> tokens = tokenize(sql);
    
    // 重新构建格式化的SQL
    return reconstruct(tokens);
}

void SQLFormatter::setIndentSize(int size)
{
    m_indentSize = size;
}

void SQLFormatter::setUppercaseKeywords(bool uppercase)
{
    m_uppercaseKeywords = uppercase;
}

QList<SQLFormatter::Token> SQLFormatter::tokenize(const QString &sql) const
{
    QList<Token> tokens;
    
    // 定义各种正则表达式模式
    QRegularExpression stringPattern(R"('(?:''|[^'])*')");  // 字符串字面量
    QRegularExpression commentPattern1(R"(--.*$)");  // 单行注释
    QRegularExpression commentPattern2(R"(/\*[\s\S]*?\*/)");  // 多行注释
    QRegularExpression numberPattern(R"(\b\d+(?:\.\d+)?\b)");  // 数字
    QRegularExpression keywordPattern(R"(\b\w+\b)");  // 关键字和标识符
    QRegularExpression operatorPattern(R"([<>!=]=?|[-+*/%,.(){}[\];])");  // 操作符
    
    QString text = sql;
    int pos = 0;
    
    while (pos < text.length()) {
        bool matched = false;
        
        // 检查字符串字面量
        QRegularExpressionMatch stringMatch = stringPattern.match(text, pos);
        if (stringMatch.hasMatch() && stringMatch.capturedStart() == pos) {
            tokens.append(Token(Token::String, stringMatch.captured(), pos));
            pos = stringMatch.capturedEnd();
            matched = true;
            continue;
        }
        
        // 检查多行注释
        QRegularExpressionMatch commentMatch2 = commentPattern2.match(text, pos);
        if (commentMatch2.hasMatch() && commentMatch2.capturedStart() == pos) {
            tokens.append(Token(Token::Comment, commentMatch2.captured(), pos));
            pos = commentMatch2.capturedEnd();
            matched = true;
            continue;
        }
        
        // 检查单行注释
        QRegularExpressionMatch commentMatch1 = commentPattern1.match(text, pos);
        if (commentMatch1.hasMatch() && commentMatch1.capturedStart() == pos) {
            tokens.append(Token(Token::Comment, commentMatch1.captured(), pos));
            pos = commentMatch1.capturedEnd();
            matched = true;
            continue;
        }
        
        // 检查数字
        QRegularExpressionMatch numberMatch = numberPattern.match(text, pos);
        if (numberMatch.hasMatch() && numberMatch.capturedStart() == pos) {
            tokens.append(Token(Token::Number, numberMatch.captured(), pos));
            pos = numberMatch.capturedEnd();
            matched = true;
            continue;
        }
        
        // 检查关键字和标识符
        QRegularExpressionMatch keywordMatch = keywordPattern.match(text, pos);
        if (keywordMatch.hasMatch() && keywordMatch.capturedStart() == pos) {
            QString word = keywordMatch.captured();
            Token::Type type = isKeyword(word) ? Token::Keyword : Token::Identifier;
            tokens.append(Token(type, word, pos));
            pos = keywordMatch.capturedEnd();
            matched = true;
            continue;
        }
        
        // 检查操作符
        QRegularExpressionMatch operatorMatch = operatorPattern.match(text, pos);
        if (operatorMatch.hasMatch() && operatorMatch.capturedStart() == pos) {
            tokens.append(Token(Token::Operator, operatorMatch.captured(), pos));
            pos = operatorMatch.capturedEnd();
            matched = true;
            continue;
        }
        
        // 如果没有匹配到任何模式，则将当前字符作为未知类型处理
        if (!matched) {
            tokens.append(Token(Token::Unknown, text.mid(pos, 1), pos));
            pos++;
        }
    }
    
    return tokens;
}

QString SQLFormatter::reconstruct(const QList<SQLFormatter::Token> &tokens) const
{
    QString result;
    int indentLevel = 0;
    bool newLine = true;
    bool afterKeyword = false;
    QString lastKeyword;
    
    for (int i = 0; i < tokens.size(); ++i) {
        const Token &token = tokens.at(i);
        
        // 处理换行和缩进
        if (newLine) {
            for (int j = 0; j < indentLevel * m_indentSize; ++j) {
                result += ' ';
            }
            newLine = false;
        }
        
        switch (token.type) {
        case Token::Keyword:
            {
                QString keyword = formatKeyword(token.value);
                
                // 特殊关键字处理
                if (keyword == "SELECT" || keyword == "INSERT" || keyword == "UPDATE" || 
                    keyword == "DELETE" || keyword == "CREATE" || keyword == "DROP" || 
                    keyword == "ALTER" || keyword == "WITH") {
                    if (i > 0) result += '\n';
                    newLine = true;
                } else if (keyword == "FROM" || keyword == "WHERE" || keyword == "GROUP" || 
                           keyword == "ORDER" || keyword == "HAVING" || keyword == "UNION" || 
                           keyword == "INTERSECT" || keyword == "EXCEPT") {
                    if (i > 0) result += '\n';
                    newLine = true;
                } else if (keyword == "BEGIN" || keyword == "CASE") {
                    indentLevel++;
                } else if (keyword == "END") {
                    if (i > 0) result += '\n';
                    indentLevel = qMax(0, indentLevel - 1);
                    newLine = true;
                }
                
                // 在某些关键字前添加额外的空行
                if ((keyword == "CREATE" || keyword == "DROP" || keyword == "ALTER" || 
                     keyword == "INSERT" || keyword == "UPDATE" || keyword == "DELETE") && 
                    i > 0 && !result.endsWith("\n\n")) {
                    result += '\n';
                }
                
                result += keyword;
                afterKeyword = true;
                lastKeyword = keyword;
            }
            break;
            
        case Token::Operator:
            {
                // 特殊操作符处理
                if (token.value == ";" && i < tokens.size() - 1) {
                    result += token.value + "\n";
                    // 在语句结束后添加额外的空行
                    if (!(i + 1 < tokens.size() && tokens.at(i + 1).type == Token::Keyword && 
                          (tokens.at(i + 1).value.toUpper() == "END" || 
                           tokens.at(i + 1).value.toUpper() == "ELSE"))) {
                        result += '\n';
                    }
                    newLine = true;
                } else if (token.value == "," || token.value == "." || 
                           token.value == "(" || token.value == ")") {
                    result += token.value;
                } else {
                    // 在操作符前后添加空格
                    result += " " + token.value + " ";
                }
                
                // 特殊处理括号
                if (token.value == "(") {
                    // 查看下一个token是否是SELECT等关键字
                    if (i + 1 < tokens.size() && tokens.at(i + 1).type == Token::Keyword) {
                        QString nextKeyword = tokens.at(i + 1).value.toUpper();
                        if (nextKeyword == "SELECT" || nextKeyword == "CASE") {
                            indentLevel++;
                            result += '\n';
                            newLine = true;
                        }
                    }
                } else if (token.value == ")") {
                    // 检查是否需要减少缩进
                    if (i + 1 < tokens.size()) {
                        Token nextToken = tokens.at(i + 1);
                        if (nextToken.type == Token::Keyword) {
                            QString nextKeyword = nextToken.value.toUpper();
                            if (nextKeyword == "SELECT" || nextKeyword == "FROM" || 
                                nextKeyword == "WHERE" || nextKeyword == "UNION") {
                                indentLevel = qMax(0, indentLevel - 1);
                                result += '\n';
                                newLine = true;
                            }
                        }
                    }
                }
            }
            break;
            
        case Token::String:
        case Token::Number:
        case Token::Identifier:
            result += token.value;
            break;
            
        case Token::Comment:
            if (token.value.startsWith("--")) {
                // 单行注释
                if (!newLine) result += ' ';
                result += token.value.trimmed();
                if (i < tokens.size() - 1) {
                    result += '\n';
                    newLine = true;
                }
            } else {
                // 多行注释
                result += token.value;
            }
            break;
            
        case Token::Whitespace:
            // 忽略原始空白，由格式化器控制
            break;
            
        case Token::Unknown:
            result += token.value;
            break;
        }
        
        // 在某些情况下添加空格
        if (i < tokens.size() - 1) {
            const Token &nextToken = tokens.at(i + 1);
            if (token.type == Token::Keyword && 
                (nextToken.type == Token::Identifier || nextToken.type == Token::String || 
                 nextToken.type == Token::Number || nextToken.type == Token::Keyword)) {
                result += ' ';
            } else if (token.type == Token::Identifier && nextToken.type == Token::Identifier) {
                result += ' ';
            }
        }
    }
    
    return result;
}

bool SQLFormatter::isKeyword(const QString &word) const
{
    return m_keywords.contains(word.toUpper());
}

QString SQLFormatter::formatKeyword(const QString &keyword) const
{
    return m_uppercaseKeywords ? keyword.toUpper() : keyword.toLower();
}